from typing import Dict, Any

"""
Convert a retrieval tool for ZhipuAI
Args:
    knowledge_id：当涉及到知识库ID时，可以前往开放平台的知识库模块进行创建或获取
    prompt_template：请求模型时的知识库模板
"""


def convert_to_retrieval_tool(knowledge_id: str, prompt_template: str = None, ) -> Dict[str, Any]:
    return ({
        "type": "retrieval",
        "retrieval": {
            "knowledge_id": knowledge_id,
            "prompt_template": prompt_template
        }
    })


"""
Convert a web_search tool for ZhipuAI
Args:
    knowledge_id: 当涉及到知识库ID时，请前往开放平台的知识库模块进行创建或获取。
    search_query: 强制搜索自定义关键内容，此时模型会根据自定义搜索关键内容返回的结果作为背景知识来回答用户发起的对话。
"""


def convert_to_web_search_tool(search_query: str = None, enable: bool = True) -> Dict[str, Any]:
    """
    将输入的查询转换为Web搜索工具配置。

    Args:
        search_query (str, optional): 要搜索的查询字符串。默认为None。
        enable (bool, optional): 是否启用Web搜索工具。默认为True。

    Returns:
        Dict[str, Any]: 包含Web搜索工具配置的字典。
    """
    return {
        "type": "web_search",
        "web_search": {
            "enable": enable,
            "search_query": search_query
        }
    }
